import matplotlib.pyplot as pl
import numpy as np
import networkx as nx
from networkx.algorithms import approximation
import itertools
import random
import collections

def draw_edge(u,v, color='g', zorder=2):
    s_x = (u[0], v[0])
    s_y = (u[1], v[1])
    pl.plot(s_x, s_y, color=color, zorder=zorder)

def distance(u,v):
    return np.sqrt((u[0] - v[0])**2 + (u[1] - v[1])**2)

def test_dominating_set(graph, ds):
    # Check if dominating set is correct
    # by definition, union of DS and N(DS) is equal to vertex set
    ds_neighbours = set()
    for dominator in ds:
        neigh = set(graph.neighbors(dominator))
        ds_neighbours = ds_neighbours.union(neigh)
    nodes = set(graph.nodes())
    union = set(ds).union(ds_neighbours)
    return nodes == union

###
# get dominating set
###

def get_dominating_set_incremental(G, reverse=False):
    # sorts nodes by 'age' and adds them to DS one by one
    points = G.nodes()
    points.sort(key=lambda x: x[2])
    root = points.pop(0)
    black_nodes = set()
    black_nodes.add(root)
    for point in points:
        # if there is black neighbor, continue
        # else, ask the eldest neighbour to be black
        black_neighbours = [n for n in G.neighbors(point) if n in black_nodes]
        if not len(black_neighbours):
            neighbors = G.neighbors(point)
            neighbors.sort(key=lambda x: x[2], reverse=reverse)
            eldest = neighbors.pop(0)
            black_nodes.add(eldest)
            #print "New black point:", eldest
            
    return black_nodes

def get_dominating_set_incremental_dummy(G):
    # sorts nodes by 'age' and adds them to DS one by one
    points = G.nodes()
    points.sort(key=lambda x: x[2])
    root = points.pop(0)
    black_nodes = set()
    black_nodes.add(root)
    for point in points:
        # if there is black neighbor, continue
        # else, ask the eldest neighbour to be black
        black_neighbours = [n for n in G.neighbors(point) if n in black_nodes]
        if not len(black_neighbours):
            neighbors = G.neighbors(point)
            random.shuffle(neighbors)
            new_point = neighbors.pop(0)
            black_nodes.add(new_point)
            #print "New black point:", eldest
            
    return black_nodes

def get_dominating_set_dummy(G):
    # Dummy algorithm of finding DS
    points = G.nodes()
    random.shuffle(points)
    root = points.pop()

    black_nodes = set()
    black_nodes.add(root)  # nodes in DS
    gray_nodes = set()  # nodes adjacent to DS
    white_nodes = set(points)  # unexplored points
    while len(white_nodes):
        root_neighbours = [n for n in G.neighbors(root) if n not in black_nodes]
        if not len(root_neighbours):
            root = random.choice(list(white_nodes))
            white_nodes.discard(root)
            black_nodes.add(root)
            continue
        random.shuffle(root_neighbours)
        root = root_neighbours.pop()
        black_nodes.add(root)
        white_nodes.discard(root)
        gray_nodes.discard(root)
        for n in root_neighbours:
            white_nodes.discard(n)
            gray_nodes.add(n)
    return black_nodes

def get_dominating_set_greedy(G, reverse=False):
    # Greedy algorithm of finding DS
    points = G.nodes()
    points.sort(key=lambda x: G.degree(x))
    root = points.pop()

    black_nodes = set()
    black_nodes.add(root)  # nodes in DS
    gray_nodes = set()  # nodes adjacent to DS
    white_nodes = set(points)  # unexplored points
    while len(white_nodes):
        root_neighbours = [n for n in G.neighbors(root) if n not in black_nodes]
        if not len(root_neighbours):
            wn = list(white_nodes)
            wn.sort(key=lambda x: G.degree(x), reverse=reverse)
            root = wn.pop()
            white_nodes.discard(root)
            black_nodes.add(root)
            continue
        root_neighbours.sort(key=lambda x: G.degree(x), reverse=reverse)
        root = root_neighbours.pop()
        black_nodes.add(root)
        white_nodes.discard(root)
        gray_nodes.discard(root)
        for n in root_neighbours:
            white_nodes.discard(n)
            gray_nodes.add(n)
    return black_nodes

def get_dominating_set_mst(G):
    mst = nx.minimum_spanning_tree(G)
    mst_DS = [v for v in mst.nodes() if mst.degree(v) > 1]  # eliminate leaves
    return mst_DS

###
#
###

def draw_dominating_set(list_DS, color='r', zorder=4):
    for DS_pair in itertools.combinations(list_DS, 2):
        u,v = DS_pair
        if distance(u,v) < radius:
            draw_edge(u,v,color=color)
    for pds in list_DS:
        x,y,i = pds
        pl.scatter(x,y, color=color, zorder=zorder)

def main():
    num_points = 500
    radius = 1.

# Generate points using 'accept-reject' algorithm
    points = [(0., 0., 1)]
# bounding box
    bbox = {
        'y+': radius,
        'y-': -radius,
        'x+': radius,
        'x-': -radius,
    }
    for i in range(2, num_points):
        while True:
            x = np.random.uniform(bbox['x-'], bbox['x+'], 1)[0]
            y = np.random.uniform(bbox['y-'], bbox['y+'], 1)[0]
            #print x,y
            new_point = (x,y,i)  # each point contains coordinates and 'age'
            is_accepted = False
            for old_point in points:
                if distance(new_point, old_point) < radius:
                    # accept
                    is_accepted = True
                    break
            if is_accepted:
                points.append(new_point)
                # Update bbox
                bbox['x-'] = min(x-1, bbox['x-'])
                bbox['x+'] = max(x+1, bbox['x+'])
                bbox['y-'] = min(y-1, bbox['y-'])
                bbox['y+'] = max(y+1, bbox['y+'])
                break

    fig = pl.figure()
    for p in points:
        circ = pl.Circle(p, radius=radius, color='#c6e6d6', alpha=0.5)
        #circ = pl.Circle(p, radius=radius, fill=False, color='b', alpha=0.5)
        fig.gca().add_artist(circ)
        x,y,i = p
        pl.scatter(x,y, zorder=3)

    G = nx.Graph()
    for point in points:
        G.add_node(point)
    for pair in itertools.combinations(points,2):
        u,v = pair
        if distance(u,v) < radius:
            G.add_edge(u,v)
            draw_edge(u, v)
    '''
    sorted_points = sorted(points, key=lambda x: x[2])
    ages = [x[2] for x in sorted_points]
    degrees = [G.degree(x) for x in sorted_points]
    pl.plot(ages, degrees)
    pl.show()
    '''
    #DS = approximation.dominating_set.min_weighted_dominating_set(G)
    #print "Size of approximated DS:", len(DS)
    dummy_DS = get_dominating_set_dummy(G)
    greedy_DS = get_dominating_set_greedy(G, reverse=True)
    inc_DS = get_dominating_set_incremental(G)
    inc_reversed_DS = get_dominating_set_incremental(G, reverse=True)
    inc_dummy_DS = get_dominating_set_incremental_dummy(G)
    mst_DS = get_dominating_set_mst(G)

    print "Size of incremental DS:", len(inc_DS)
    print "Size of incremental reversed DS:", len(inc_reversed_DS)
    #print "Is incremental DS correct?", "Yes!" if test_dominating_set(G, inc_DS) else "No..."
    print "Size of dummy incremental DS:", len(inc_dummy_DS)
    #print "Is dummy incremental DS correct?", "Yes!" if test_dominating_set(G, inc_dummy_DS) else "No..."
    print "Size of greedy DS:", len(greedy_DS)
    #print "Is greedy DS correct?", "Yes!" if test_dominating_set(G, greedy_DS) else "No..."
    print "Size of dummy DS:", len(dummy_DS)
    #print "Is dummy DS correct?", "Yes!" if test_dominating_set(G, dummy_DS) else "No..."
    print "Size of MST DS: ", len(mst_DS)
    degrees_distr = collections.Counter(G.degree().values())
#for deg in sorted(degrees_distr.keys()):
#    print deg, degrees_distr[deg]

    '''
    pl.axis('equal')
    draw_dominating_set(list(mst_DS), color='y')
    draw_dominating_set(list(inc_DS), color='r')
    pl.show()
    '''

if __name__ == '__main__': main()
